import warnings

import joblib
from sklearn import metrics

warnings.filterwarnings("ignore")

import pandas as pd
from imblearn.over_sampling import SMOTE
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from xgboost import XGBClassifier

# 数据加载
df = pd.read_csv('healthcare-dataset-stroke-data.csv')
df_enc = df.iloc[:, 1:]
# 数据编码转换
df_enc['gender'] = df_enc['gender'].replace(
    {
        'Male': 0,
        'Female': 1,
        'Other': 2
    }
)
df_enc['ever_married'] = df_enc['ever_married'].replace(
    {
        'No': 0,
        'Yes': 1
    }
)

df_enc['Residence_type'] = df_enc['Residence_type'].replace(
    {
        'Rural': 0,
        'Urban': 1
    }
)
df_enc['work_type'] = df_enc['work_type'].replace(
    {
        'children': 0,
        'Govt_job': 1,
        'Never_worked': 2,
        'Private': 3,
        'Self-employed': 4
    }
)
df_enc['smoking_status'] = df_enc['smoking_status'].replace(
    {
        'formerly smoked': 0,
        'never smoked': 1,
        'smokes': 2,
        'Unknown': 3
    }
)

df_enc['bmi'] = df_enc['bmi'].fillna(0)
# 预测值
Y = df_enc['stroke']
# 特征
X = df_enc.drop(['stroke'], axis=1)
# 数据划分
x_train, x_test, y_train, y_test = train_test_split(X, Y, train_size=0.8, random_state=6, shuffle=True)
# 数据标准化，注意，train和test要分开标准化，否则有数据泄露风险
scaler_train = StandardScaler()
scaler_train.fit(x_train)
scaler_test = StandardScaler()
scaler_test.fit(x_train)
x_train = scaler_train.transform(x_train)
x_test = scaler_test.transform(x_test)
# 数据不平衡处理,采样后比例为1：1，注意，只对训练集做数据不平衡处理
sm = SMOTE(random_state=6, sampling_strategy=1.0)
x_train, y_train = sm.fit_resample(x_train, y_train)
# 模型训练
xgbs = XGBClassifier(random_state=6)
xgbs.fit(x_train, y_train)
test_predict = xgbs.predict(x_test)

print(" Accuracy:",
      format(metrics.accuracy_score(y_test, test_predict), '.4f'),
      "Precision:", format(metrics.precision_score(y_test, test_predict), '.4f'),
      "Recall:", format(metrics.recall_score(y_test, test_predict), '.4f'),
      "F1:", format(metrics.f1_score(y_test, test_predict), '.4f'),
      "AUC:", format(metrics.roc_auc_score(y_test, test_predict), '.4f'))
# 模型生成
joblib.dump(xgbs, "smote_xgb.pkl")
